# 实现混合高斯提取视频背景
import cv2
import numpy as np
import os
import math

# [170:270,350:480],由于计算很慢，取部分图像计算观察
tx,ty,tw,th = 150,350,130,130

class my_BackgroundSubtractor():

    def __init__(self,alpha=0.05,k=3,threshold_inx=2,w=0.1,s=10):
        '''
        构造函数
        @alpha  参数更新的学习率，默认为0.05
        @k      混合高斯模型的数目，默认为3个
        @threshold_inx  背景判断的阈值，默认为2,取前2个模型
        @w      新建高斯模型的权值
        @s      新建高斯模型的方差
        '''
        self.alpha_ = alpha
        self.k_ = k
        self.threshold_inx = threshold_inx
        self.w_ = w
        self.s_ = s
        self.lastImg = np.array([0])
        self.param = np.array([0])

    def apply(self,img):
        '''
        对图像进行背景提取
        @img    图像帧
        @return 二值化前景图像
        '''
        # 转成灰度图
        self.img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)[tx:tx+tw,ty:ty+th]
        # 创建背景和前景图,背景存二值化的值，0和1
        self.backImg = self.img.copy()
        self.frontImg = np.zeros_like(self.img,dtype='float')
        # 存储尺寸
        w,h = self.img.shape
        # 创建参数容器 4个维度，前两个对应的像素坐标，第三个为第高斯分量的个数，第4个为高斯分量的权值，均值，方差
        # 初始化第一个模型
        if not (self.param.any() > 0):
            self.param = np.zeros((w,h,self.k_,3),dtype='float')
            a = np.concatenate((np.ones((w,h,1))*1,self.img[:,:,np.newaxis]),axis=2)
            self.param[:,:,0] = np.concatenate((a,np.ones((w,h,1))*self.s_),axis=2)
            self.lastImg = self.img.copy()
            return self.img
        # 遍历图像
        for i in range(w):
            for j in range(h):
                v = self.img[i,j]
                idx = self.updateModel(v,(i,j))
                # 如果是前景图像，则使用权值最大的高斯模型的均值代替
                if idx == -1 or idx > self.threshold_inx:
                    self.backImg[i,j] = self.param[i,j,0,1]
                    self.frontImg[i,j] = 1.0
                else:
                    self.frontImg[i,j] = 0.0

        # 用np操作
        # # 存储当前灰度落在哪个高斯模型的掩码
        # isBack = np.zeros((w,h,self.k_))
        # isBack[abs(self.param[:,:,:,1] - self.img[:,:,np.newaxis]) < self.param[:,:,:,2]]=1
        # isBack[abs(self.param[:,:,:,1] - self.img[:,:,np.newaxis]) >= self.param[:,:,:,2]]=0
        # # 没有在范围内，则新建一个模型，替换最后一个
        # # 如果在设定范围内，则设置背景前景图片
        # for x in range(len(isBack)):
        #     for y in range(len(isBack[x])):
        #         if(isBack[x,y].all() == 0):
        #             self.param[x,y,2] = np.array([self.w_,self.img[x,y],self.s_])
        #         if(isBack[x,y,:self.threshold_inx].any() == 1):
        #             self.backImg[x,y] = self.lastImg[x,y]
        #             self.frontImg[x,y] = 1.0
        #         else:
        #             self.frontImg[x,y] = 0.0
      
        # # 更新权值
        # self.param[:,:,:,0] = (1 - self.alpha_) * self.param[:,:,:,0] + isBack*self.alpha_
        # # 计算ρ
        # rho = self.alpha_ / self.param[:,:,:,0]
        # # 更新mu
        # self.param[:,:,:,1] = (1 - rho) * self.param[:,:,:,1] + rho * self.img[:,:,np.newaxis]
        # # 更新sigma
        # self.param[:,:,:,2] = (1-rho) * (self.param[:,:,:,2]**2) + rho * (self.img[:,:,np.newaxis] - self.param[:,:,:,1])*(self.img[:,:,np.newaxis] - self.param[:,:,:,1])
        # self.param[:,:,:,2] = self.param[:,:,:,2] ** 0.5
        # # 排序
        # for x in self.param:
        #     for y in x:
        #         y = y[np.argsort(y[:,0])[::-1]]

        # 存储上一张前景图像
        self.lastImg = self.backImg
        # 返回前景
        return self.frontImg

    def getBackGroundImg(self):
        '''
        得到背景图像
        @return 背景图像
        '''
        return self.backImg

    def gaussCul(self,mu,sigma,v):
        '''
        高斯函数，输入参数，计算响应的高斯函数值
        @mu     均值
        @sigma  方差
        @return 计算结果
        @v      要计算的值
        @return 如果高斯模型还没建立，则返回0.0
        '''
        if ((abs(mu) <= 0.01) and (abs(sigma) <= 0.01)):
            return 0.0
        return math.exp(-((v-mu)**2)/(2*(sigma**2)))/(2.5066*sigma)

    def updateModel(self,v,axi=(0,0)):
        '''
        模型参数更新
        @v      当前像素的灰度值
        @axi  需要跟新的位置
        @return 返回高斯模型的索引，如果没有落在任何一个中，则返回-1
        '''
        x,y = axi
        guassModel = self.param[x,y]
        # 判断灰度落在哪个高斯模型中，没有在任何一个中，则取-1
        index = -1
        for idx,(_,mu,sigma) in enumerate(guassModel):
            if abs(v-mu)<2.5*sigma:
                index = idx
                break
        if index == -1:
            #移除最小权值的模型，并新建高斯模型 
            newGuass = np.array([self.w_,v,self.s_])
            guassModel[-1] = newGuass
        else:
            # 更新高斯模型
            for idx,(w,mu,sigma) in enumerate(guassModel):
                # 计算ρ,老师讲的方法要计算高斯函数，比较慢
                rho = self.alpha_*self.gaussCul(mu,sigma,v)
                # 如果这个高斯模型还没建立，则计算下一个
                if w == 0.0:
                    continue
                if idx == index:
                    guassModel[idx,0] = (1 - self.alpha_) * w + self.alpha_ 
                else:
                    guassModel[idx,0] = (1 - self.alpha_) * w
                # 换成网友提供的
                # rho = self.alpha_ / guassModel[idx,0]
                r = 0.005
                guassModel[idx,1] = (1 - rho) * guassModel[idx,1] + rho * v
                guassModel[idx,2] = (1 - rho) * (sigma ** 2) + rho * ((v-guassModel[idx,1])**2)
                guassModel[idx,2] = guassModel[idx,2] ** 0.5
        # 按权值排序
        self.param[x,y] = guassModel[np.argsort(guassModel[:,0])[::-1]]
        return index
        

cap = cv2.VideoCapture(os.path.dirname(__file__)+'/vtest.avi')
mybgs = my_BackgroundSubtractor(alpha=0.05,k=4,threshold_inx=2,w=0.9)

count = 0
while True:
    ret,frame = cap.read()
    # 没有图像了就退出 
    if not ret:
        break
    count += 1
    if ((count % 5) != 0):
        continue
    fgimg = mybgs.apply(frame)
    bg = mybgs.getBackGroundImg()

    cv2.putText(frame,str(count),(ty - 60,tx),cv2.FONT_HERSHEY_PLAIN,1.5,color=[0,0,255])
    cv2.rectangle(frame,(ty,tx),(ty+th,tx+tw),color=[0,0,255])

    cv2.imshow('oring',frame)
    cv2.imshow('front',fgimg)
    cv2.imshow('back',bg)
    c = cv2.waitKey(20)
    if c == -1:
        pass
    elif chr(c) == ' ':
        cv2.waitKey()
    print(count)

cv2.destroyAllWindows()

